#include "Segmentation.h"
#include "Matrix.h"
#include <algorithm>
#include <stdio.h>

using namespace std;

#define PI 3.14159265358979323846

Segmentation::Segmentation(int heightValue, int widthValue, int depthValue, int numberOfImagesValue, int numberOfRegionsValue, LabelID* initialPointsValue, ImageType* imagesValue){
	height = heightValue;
	width = widthValue;
	depth = depthValue;
	numberOfImages = numberOfImagesValue;
	numberOfRegions = numberOfRegionsValue;
	initialPoints = initialPointsValue;
	images = imagesValue;
	mask = NULL;
	mean = NULL;
	variance = NULL;
	covariance = NULL;
	dataCost = NULL;
	smoothnessCost = NULL;
	graph = NULL;

	bounds.minX = 0;
	bounds.maxX = height-1;
	bounds.minY = 0;
	bounds.maxY = width-1;
}

Segmentation::~Segmentation() {
	if (graph != NULL)
		delete graph;
	if (mean != NULL)
		delete [] mean;
	if (variance != NULL)
		delete [] variance;
	if (covariance != NULL)
		delete [] covariance;
	if (dataCost != NULL)
		delete [] dataCost;
	if (smoothnessCost != NULL)
		delete [] smoothnessCost;
}

/*
Calculate Mean and Covariance in order to obtain data cost based on
a Multivariate normal distribution.
*/
void Segmentation::calculateMeanAndCovariance()
{
	int* contador =	new int[numberOfRegions];	//contador para cada regionMean

	for(int i=0;i<numberOfRegions;i++)
		contador[i] = 0;

	if (mean != NULL){
		delete [] mean;
		delete [] covariance;
		delete [] variance;
	}
	mean = new double[numberOfRegions*numberOfImages];
	variance = new double[numberOfRegions*numberOfImages];
	covariance = new double[numberOfRegions*numberOfImages*numberOfImages];

	int** dataPerRegion = new int*[numberOfRegions*numberOfImages];

	for(int i=0;i<numberOfRegions*numberOfImages;i++){

		mean[i] = 0;
		variance[i] = 0;
		dataPerRegion[i] = new int[height*width*depth];

		for(int j=0;j<height*width*depth;j++)
			dataPerRegion[i][j] = 0;
	}

	int iterator = 0;
	int value = 0;

	for(int i=0;i<height*width*depth; i++){

		if( (int)initialPoints[i] != -1 && (mask == NULL || (int) mask[i] > 0)){
			for(int y=0;y<numberOfImages;y++){

				iterator = (int) initialPoints[i] + y*numberOfRegions;

				value = (int) images[i + y*height*width*depth];

				mean[iterator] += value;
				
				dataPerRegion[iterator][contador[(int)initialPoints[i]]] = value;

			}
			contador[(int)initialPoints[i]]++;
		}

	}

	for(int j=0;j<numberOfRegions; j++){
		if(printFlag == 1)
			printf("Region %d \n",j);

		for(int i=0;i<numberOfImages;i++){
			if(contador[j] == 0)
				contador[j] = 1;

			iterator = j + i*numberOfRegions;
			mean[iterator] = mean[iterator]/contador[j];

			if(printFlag == 1)
				printf("Im %d %f -",i,mean[j + i*numberOfRegions]);
		}
		if(printFlag == 1)
			printf("\n");
	}

	int regionMean = 0;
	double cov = 0.0;
	for(int i=0;i<numberOfRegions; i++){

		if(printFlag == 1)
			printf("Region %d \n",i);

		for(int x=0;x<numberOfImages;x++){

			for(int y=0;y<numberOfImages;y++){

				iterator = y + x*numberOfImages;
				regionMean = i*numberOfImages*numberOfImages;
				covariance[iterator + regionMean] = 0;

				if (contador[i] == 1)
					contador[i]++;

				for(int cont=0;cont<contador[i];cont++){

					cov = dataPerRegion[i + x*numberOfRegions][cont]-mean[i + x*numberOfRegions];
					cov = cov * (dataPerRegion[i + y*numberOfRegions][cont]-mean[i + y*numberOfRegions]);
					covariance[iterator + regionMean] += cov/(contador[i]-1);
				}


				if(x == y){
					cov = covariance[iterator + regionMean];
					covariance[iterator + regionMean] = cov + pow(2.220446049250313,-16);
				}

				if(printFlag == 1)
					if(y == 0 && x == 0)
						printf("[ %f ",covariance[iterator + regionMean]);
					else if ( y == 0 )
						printf(" %f ",covariance[iterator + regionMean]);
					else
						printf(", %f ",covariance[iterator + regionMean]);
        		}

			if(printFlag == 1)
				printf("; ");

			for(int cont=0;cont<contador[i];cont++){
				variance[i + x*numberOfRegions] += pow(dataPerRegion[i + x*numberOfRegions][cont]-mean[i + x*numberOfRegions],2);
			}

			if (numberOfRegions > 1)
				variance[i + x*numberOfRegions] = variance[i + x*numberOfRegions]/(numberOfRegions - 1);
        	}

	        if(printFlag == 1)
				printf("] \n ");
	}

	for(int i=0;i<numberOfRegions*numberOfImages;i++){
		delete [] dataPerRegion[i];
	}
	delete [] dataPerRegion;
	delete [] contador;

}

/*
	Export region Mean data to python
*/
void Segmentation::exportMean(int region, double* exportMean){
	if (mean == NULL)
		calculateMeanAndCovariance();
	for(int i=0;i<numberOfImages;i++)
		exportMean[i] = mean[region + i*numberOfRegions];
}

/*
	Export region Covariance data to python
*/
void Segmentation::exportCovariance(int region, double* exportCov){
	if (mean == NULL)
		calculateMeanAndCovariance();
	for(int i=0;i<numberOfImages*numberOfImages;i++)
		exportCov[i] = covariance[i + region*numberOfImages*numberOfImages];

}

/*
	Export region Variance data to python
*/
void Segmentation::exportVariance(int region, double* exportVar){
	if (mean == NULL)
		calculateMeanAndCovariance();
	for(int i=0;i<numberOfImages;i++)
		exportVar[i] = variance[region + i*numberOfRegions];
}

/*
	Gets regions curtosis
*/

void Segmentation::getCurtosis(int region,double* returnValue){
	if (mean == NULL)
		calculateMeanAndCovariance();

	int contador = 0;

	double* dividend = new double[numberOfImages];
	double* divisor = new double[numberOfImages];

	for( int i=0 ; i < numberOfImages ; i++){
		dividend[i] = 0;
		divisor[i] = 0;
	}

	for( int i=0 ; i < height*width*depth ; i++){
		if( (int)initialPoints[i] == region && (mask == NULL || (int) mask[i] > 0)){
			for( int y=0 ; y < numberOfImages ; y++ ){
				dividend[y] += pow((double) images[i + y*height*width*depth] - mean[y*numberOfRegions + region],4);
				divisor[y] += pow((double) images[i + y*height*width*depth] - mean[y*numberOfRegions + region],2);
			}
			contador++;
		}
	}

	for( int y=0 ; y < numberOfImages ; y++ ){
		dividend[y] = dividend[y] * contador;
		divisor[y] = pow(divisor[y],2);
		returnValue[y] = (dividend[y] / divisor[y]) - 3;
	}

	delete[] dividend;
	delete[] divisor;
}

/*
	Set mask and obtain bounds for it
*/
void Segmentation::setMask(LabelID* newMask){

	mask = newMask;
	
	bounds.minX = height;
	bounds.maxX = 0;
	bounds.minY = width;
	bounds.maxY = 0;

	for(int i=0; i<depth; i++){
		for(int x=0; x<height; x++){
			for(int y=0; y<width; y++){
				if( (int) mask[y + width*x + height*width*i] > 0){
					bounds.minX = min(bounds.minX,x);
					bounds.maxX = max(bounds.maxX,x);
					bounds.minY = min(bounds.minY,y);
					bounds.maxY = max(bounds.maxY,y);
				}
			}
		}
	}
	if(printFlag == 1)
		printf(" Bounds = [ %d , %d , %d , %d ] \n",bounds.minX, bounds.maxX, bounds.minY, bounds.maxY);
}


void Segmentation::setBounds(int minXValue, int maxXValue, int minYValue, int maxYValue){
	bounds.minX = minXValue;
	bounds.maxX = maxXValue;
	bounds.minY = minYValue;
	bounds.maxY = maxYValue;
}

/*
	Returns 1 if passed coordante is inside the mask boundaries.
*/
int Segmentation::insideBounds(int x, int y){
	int ret = 0;
	if(x >= bounds.minX && x <= bounds.maxX && y >= bounds.minY && y <= bounds.maxY)
		ret = 1;
	return ret;
}

/*
	Set smoothFactor in order to run segmentation
*/
void Segmentation::setSmoothFactor(double factor){
	if (smoothnessCost!= NULL)
		delete [] smoothnessCost;

	smoothnessCost = new Graph::captype[numberOfRegions*numberOfRegions];

	for(int i=0;i<numberOfRegions;i++)
		for(int j=0;j<numberOfRegions;j++)
			if(i != j)
				smoothnessCost[i + j*numberOfRegions] = factor;
			else
				smoothnessCost[i + j*numberOfRegions] = 0;
}

/*
	Obtain dataCost from Multivariate normal distribution
*/
void Segmentation::calculateDataCost(){
	int maskHeight = bounds.maxX - bounds.minX + 1;
	int maskWidth = bounds.maxY - bounds.minY + 1;

	if(dataCost == NULL)
		dataCost = new double[maskHeight*maskWidth*depth*numberOfRegions];

	calculateMeanAndCovariance();

	Matrix** meanMatrix = new Matrix*[numberOfRegions];
	Matrix** covarianceMatrix = new Matrix*[numberOfRegions];

	double* regionMean = new double[numberOfImages];
	double* regionCovariance = new double[numberOfImages*numberOfImages];

	for(int i=0;i<numberOfRegions;i++){
		meanMatrix[i] = new Matrix(numberOfImages);

		exportMean(i,regionMean);

		meanMatrix[i]->set(regionMean);

		covarianceMatrix[i] = new Matrix(numberOfImages,numberOfImages);

		exportCovariance(i,regionCovariance);

		covarianceMatrix[i]->set(regionCovariance);
	}

	Matrix* multiplicando = new Matrix(1,numberOfImages);
	Matrix* escalar = new Matrix();
	Matrix* pixel = new Matrix(numberOfImages);
	Matrix* rest = new Matrix(numberOfImages);
	Matrix* transRest = new Matrix(1,numberOfImages);

	double value = 0.0;
	double plus = 0.0;

	Matrix** inverses = new Matrix*[numberOfRegions];
	double* determinants = new double[numberOfRegions];

	for(int i=0;i<numberOfRegions;i++){

		inverses[i] = new Matrix(numberOfImages,numberOfImages);
		determinants[i] = inv(*covarianceMatrix[i], *inverses[i]);

		if(printFlag == 1)
			printf("Inversa de %d con determinant %f \n",i,determinants[i]);

		if(printFlag == 1)
			inverses[i]->show();
	}

	for(int i=0; i<depth; i++){
		for(int x=0; x<maskHeight; x++){
			for(int y=0; y<maskWidth; y++){
				if(mask != NULL && mask[(y+bounds.minY) + width*(x+bounds.minX) + height*width*i] == 0)
					for(int p=0;p<numberOfImages;p++)
						pixel->set(p,0.0);
				else
					for(int p=0;p<numberOfImages;p++)
						pixel->set(p, images[(y+bounds.minY) + width*(x+bounds.minX) + height*width*i + p*height*width*depth]);

				for(int l=0; l<numberOfRegions; l++){
					rest->substract(*pixel,*meanMatrix[l]);

					transRest->transpose(*rest);

					multiplicando->multiply(*transRest,*inverses[l]);

					escalar->multiply(*multiplicando,*rest);

					value = (*escalar)(0,0)/2;

					plus = sqrt(fabs(determinants[l]))*pow(2*PI,numberOfImages/2);

					if (value + log(plus) >= 0)
						value = value + log(plus);
					dataCost[(y + maskWidth*x + maskWidth*maskHeight*i)*numberOfRegions + l] = value;

				}
			}
		}
	}

	//chequear bien delets
	delete rest;
	delete transRest;
	delete pixel;
	delete escalar;
	delete multiplicando;
	for(int i=0;i<numberOfRegions;i++){
			delete meanMatrix[i];
			delete covarianceMatrix[i];
			delete inverses[i];
	}
	delete[] determinants;
	delete[] inverses;
	delete[] covarianceMatrix;
	delete[] meanMatrix;
	delete[] regionMean;
	delete[] regionCovariance;
}

/*
	Export dataCost to python
*/
void Segmentation::exportDataCost(double* pythonDataCost){

	int maskHeight = bounds.maxX - bounds.minX + 1;
	int maskWidth = bounds.maxY - bounds.minY + 1;

	if(dataCost == NULL)
		calculateDataCost();
	
	for(int i=0;i<maskHeight*maskWidth*depth*numberOfRegions;i++){
		pythonDataCost[i] = dataCost[i];
	}
}

/*
	Expand regions without calling the graph cut segmentation	
*/
void Segmentation::expandRegions(int iterations){

	int maskHeight = bounds.maxX - bounds.minX + 1;
	int maskWidth = bounds.maxY - bounds.minY + 1;

	if (dataCost == NULL)
		calculateDataCost();

	double prevEnergy = 0;
	double energy = 0;

	double minEnergy = 0;
	int region = 0;
	
	bool keepGoing = true;

	while(keepGoing){

	   energy = 0;

	   for(int i=0; i<depth; i++){
	      for(int x=0; x<height; x++){
	         for(int y=0; y<width; y++){
	            if(insideBounds(x,y) == 1){
	               minEnergy = dataCost[(y-bounds.minY + maskWidth*(x-bounds.minX) + maskWidth*maskHeight*i)*numberOfRegions];
	               region = 0;
	               for(int r=1;r<numberOfRegions;r++){
	                  if (minEnergy > dataCost[(y-bounds.minY + maskWidth*(x-bounds.minX) + maskWidth*maskHeight*i)*numberOfRegions + r]){
	                     region = r;
	                     minEnergy = dataCost[(y-bounds.minY + maskWidth*(x-bounds.minX) + maskWidth*maskHeight*i)*numberOfRegions + r];
	                  }
	               }
	               energy = energy + minEnergy;
	               initialPoints[y + x*width + i*width*height] = region;
	            } else{
	               initialPoints[y + x*width + i*width*height] = -1;
	            }
	         }
	      }
	   }

	   keepGoing = ( prevEnergy == 0 || prevEnergy > energy ) && ( iterations == 0 || iterations > 1 );
	   if (iterations > 1) 
	      iterations--;
	   prevEnergy = energy;

	   printf("Energy = %f \n", energy);

	   if (keepGoing)
	      calculateDataCost();
	}
}

/*
	Set initialPoints
*/
void Segmentation::setInitialPoints(LabelID* points){
	initialPoints = points;
}

/*
	Export initialPoints
*/
void Segmentation::exportInitialPoints(LabelID* points){

	assert(initialPoints != NULL);

	for(int i=0;i<height*width*depth;i++){
		points[i] = initialPoints[i];
	}
}

/*
	Update InitialPoints from points inside the mask
*/
void Segmentation::updateInitialPoints(LabelID *points){
	int maskHeight = bounds.maxX - bounds.minX + 1;
	int maskWidth = bounds.maxY - bounds.minY + 1;

	for(int x=0;x<maskHeight;x++){
		for(int y=0;y<maskWidth;y++){
			for(int z=0;z<depth;z++){
				initialPoints[y + bounds.minY + (x + bounds.minX)*width + width*height*z] = points[y + x*maskWidth + maskWidth*maskHeight*z];
			}
		}
	}
}

/*
	Run expansion algorithm iterations times
*/
void Segmentation::expansion(int iterations){
	initGraph();

	if (dataCost == NULL)
		calculateDataCost();

	graph->setDataCost(dataCost);

	int maskHeight = bounds.maxX - bounds.minX + 1;
	int maskWidth = bounds.maxY - bounds.minY + 1;

	double past_energy = 0;
	double energy = 0;

	bool keepGoing = true;
	int i = 0;

	LabelID* pixels = NULL;

	while(keepGoing){
		i++;
		past_energy = energy;

		energy = graph->expansion(2);
		printf("Iteration - %d - ",i);
		printf("Energy %f\n",energy);

		if(pixels == NULL)
			pixels = new LabelID[maskWidth*maskHeight*depth];

		graph->ExportLabels(pixels);

		updateInitialPoints(pixels);
		if((past_energy != 0 && past_energy <= energy) || (iterations != 0 && i >= iterations))
			keepGoing = false;

		if (keepGoing){

			calculateDataCost();
			graph->reSetDataCost(dataCost);
		}

	}
	delete [] pixels;
}

/*
	Configure graph
*/
void Segmentation::initGraph(){
	if (graph != NULL)
		delete graph;

	int maskHeight = bounds.maxX - bounds.minX + 1;
	int maskWidth = bounds.maxY - bounds.minY + 1;

	graph = new GCoptimizationGeneralGraph(maskHeight*maskWidth*depth, numberOfRegions);
	graph->setSmoothCost(smoothnessCost);
	setFullNeighbours();
}

/*
	Return 1 if a node is inside the mask in order to be connected in the graph structure
*/
int Segmentation::isConnected(int x, int y, int z){
	int ret = 0;
	if(mask == NULL || mask[y + bounds.minY + width*(x + bounds.minX) + width*height*z] != 0)
		ret = 1;
	return ret;
}

/*
	Set Full Neighbour configuration
*/
void Segmentation::setFullNeighbours(){
	int x,y,z;

	double weight;

	int node = 0;
	int neighbours = 0;

	int maskHeight = bounds.maxX - bounds.minX + 1;
	int maskWidth = bounds.maxY - bounds.minY + 1;

	z = 0;
	while(z < depth){
	   x = 0;
	   while (x < maskHeight) {
	      y = 0;
	      while (y < maskWidth) {
	         neighbours = 0;
	         if (isConnected(x,y,z) == 1){
	            node = y + maskWidth*x + maskWidth*maskHeight*z;

	            // SAME SLIDE NEIGHBOURS
	            neighbours = neighbours + setNeighboursInSlide(x,y,z,z);
	
	            // NEXT SLIDE NEIGHBOURS
	            if(z < depth-1){
	               neighbours = neighbours + setNeighboursInSlide(x,y,z,z+1);

	               //Same pixel, next slide (x,y,z+1)
	               if(isConnected(x,y,z+1) == 1){
	                  graph->setNeighbors(node, node + maskWidth*maskHeight, sliceDistance);
	                  neighbours++;
	               }
	
	               //Prev column, Same row, next slide (x,y-1,z+1)
	               if (y > 0 && isConnected(x,y-1,z+1) == 1){
	                  graph->setNeighbors(node, node - 1 + maskWidth*maskHeight,sqrt(1+pow(sliceDistance,2)));
	                  neighbours++;
	               }
	
	               //Same column, Prev row, next slide (x-1,y,z+1)
	               if( x > 0 && isConnected(x-1,y,z+1) == 1){
	                  graph->setNeighbors(node, node - maskWidth + maskWidth*maskHeight,sqrt(1+ pow(sliceDistance,2)));
	                  neighbours++;
	               }
	
	               //Prev column, Prev row, next slide (x-1,y-1,z+1)
	               weight = sqrt(2.0);
	               if(sliceDistance != 0)
	                  weight = sqrt( pow(weight,2) + pow(sliceDistance,2));
	               if ( x > 0 && y > 0 && isConnected(x-1,y-1,z+1) == 1){
	                  graph->setNeighbors(node, node - 1 - maskWidth + maskWidth*maskHeight,weight);
	                  neighbours++;
	               }
	
	               //Next column, Prev row, next slide (x-1,y+1,z+1)
	               if ( y < maskWidth - 1 && x > 0 && isConnected(x-1,y+1,z+1) == 1){
	                  graph->setNeighbors(node, node + 1 - maskWidth + maskWidth*maskHeight,weight);
	                  neighbours++;
	               }
	            }
	         }
	         y++;
	      }
	      x++;
	   }
	   z++;
	}
}

int Segmentation::setNeighboursInSlide(int x, int y, int zFrom, int zTo){

	int maskHeight = bounds.maxX - bounds.minX + 1;
	int maskWidth = bounds.maxY - bounds.minY + 1;

	int numberNeighbours = 0;

	int nodeFrom = y + maskWidth*x + maskWidth*maskHeight*zFrom;

	//(x,y+1,zTo)
	int nodeTo = (y + 1) + maskWidth*x + maskWidth*maskHeight*zTo;
	double weight = 0;

	if ( y < maskWidth - 1 && isConnected(x,y+1,zTo) == 1){
		graph->setNeighbors(nodeFrom, nodeTo, sqrt(1+sliceDistance));
		numberNeighbours++;
	}

	//(x+1,y,zTo)
	nodeTo = y + maskWidth*(x+1) + maskWidth*maskHeight*zTo;


	if (x < maskHeight -1){

		if (isConnected(x+1,y,zTo) == 1) {
			graph->setNeighbors(nodeFrom, nodeTo ,sqrt(1+pow(sliceDistance,2)));
			numberNeighbours++;
		}

		weight = sqrt(2.0);
		if(sliceDistance != 0)
			weight = sqrt( pow(weight,2) + pow(sliceDistance,2));

		//(x+1,y+1,zTo)
		nodeTo++;
		if ( y < maskWidth - 1 && isConnected(x+1,y+1,zTo) == 1){
			graph->setNeighbors(nodeFrom, nodeTo, weight);
			numberNeighbours++;
		}

		//(x+1,y-1,zTo)
        	nodeTo = nodeTo-2;
		if ( y > 0 && isConnected(x+1,y-1,zTo) == 1){
			graph->setNeighbors(nodeFrom, nodeTo, weight);
			numberNeighbours++;
		}
	}

	return numberNeighbours;
}
